import os

from utils.data_load import dataload
from utils.feature_engineering import feature_engineering
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, roc_curve, accuracy_score
import joblib
import matplotlib.pyplot as plt
import seaborn as sns


def test_model():
    # 1. 数据加载
    X_train, X_test, y_train, y_test = dataload()

    # 2. 特征工程
    X_train, X_test, y_train = feature_engineering(X_train, X_test, y_train)

    # 3. 加载模型
    model = joblib.load("../model/logistic_regression_model.pkl")

    # 4. 模型预测
    y_pred_train = model.predict(X_train)
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]
    accuracy_train = accuracy_score(y_train, y_pred_train)
    accuracy_test = accuracy_score(y_test, y_pred)
    print('===== 训练集和测试集模型预测准确率 =====')
    print('Training Accuracy: %.2f%%' % (accuracy_train * 100))
    print('Testing Accuracy: %.2f%%' % (accuracy_test * 100))

    # 5. 评估结果
    print("===== Classification Report =====")
    print(classification_report(y_test, y_pred))

    cm = confusion_matrix(y_test, y_pred)
    print("\n===== Confusion Matrix =====")
    print(cm)

    auc = roc_auc_score(y_test, y_proba)
    print(f"\nAUC: {auc:.4f}")

    # 6. 可视化
    plt.figure(figsize=(12, 5))

    # 混淆矩阵热力图
    plt.subplot(1, 2, 1)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
    plt.title('Confusion Matrix')
    plt.xlabel('Predicted')
    plt.ylabel('True')

    # ROC 曲线
    fpr, tpr, _ = roc_curve(y_test, y_proba)
    plt.subplot(1, 2, 2)
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'AUC = {auc:.3f}')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.title('ROC Curve')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(loc="lower right")

    plt.tight_layout()
    os.makedirs("../fig", exist_ok=True)
    fig_path = "../fig/logistic_regression_eval.png"
    plt.savefig(fig_path, dpi=300, bbox_inches='tight')


if __name__ == "__main__":
    test_model()
